import queue
import sys
import pyaudio
import torch
import numpy as np
import librosa
import torchaudio
import random
import tkinter as tk
from tkinter import scrolledtext, Button, Label, Entry
import threading

sys.path.append('third_party/Matcha-TTS')
from cosyvoice.cli.cosyvoice import CosyVoice2

#import os
#os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "backend:cudaMallocAsync"  # 异步分配 设置必须在PyTorch初始化（即任何CUDA操作）之前完成

# 严格匹配官方音频处理参数
max_val = 0.8
TOP_DB = 60
HOP_LENGTH = 220
WIN_LENGTH = 440
SAMPLE_RATE = 24000  # CosyVoice的标准采样率

# 初始化PyAudio音频流
p = pyaudio.PyAudio()
stream = p.open(
    format=pyaudio.paFloat32,
    channels=1,
    rate=SAMPLE_RATE,
    output=True,
    frames_per_buffer=512
)

# 初始化语音合成模型 - 严格匹配官方参数
cosyvoice = CosyVoice2(
    'pretrained_models/CosyVoice2-0.5B',
    load_jit=False,
    load_trt=False,
    fp16=True # 启用半精度
)

torch.backends.cudnn.benchmark = True  # 自动选择最优卷积算法
torch.set_num_threads(4)  # 限制CPU线程避免资源争抢

# 官方精确复制版音频处理功能
def official_audio_processing(audio_data, sample_rate=SAMPLE_RATE):
    """
    完全复制CosyVoice官方demo的音频处理流程
    """
    # 1. 转换为PyTorch张量
    if isinstance(audio_data, np.ndarray):
        speech = torch.from_numpy(audio_data).float()
    else:
        speech = audio_data.float()
    
    # 2. 单声道处理
    if len(speech.shape) > 1:
        if speech.shape[0] > 1:
            speech = speech.mean(0, keepdim=True)
        speech = speech.squeeze(0)
    
    # 3. 精确的静音修剪 (完全匹配官方实现)
    speech = speech.unsqueeze(0)
    speech, _ = librosa.effects.trim(
        speech.squeeze(0).numpy(),
        top_db=TOP_DB,
        frame_length=WIN_LENGTH,
        hop_length=HOP_LENGTH
    )
    speech = torch.tensor(speech).float().unsqueeze(0)
    
    # 4. 精确的音量归一化
    max_amp = speech.abs().max()
    if max_amp > max_val:
        speech = speech / max_amp * max_val
    
    # 5. 添加结尾静音缓冲 (官方流程)
    silence_duration = int(0.2 * SAMPLE_RATE)
    speech = torch.concat(
        [speech, torch.zeros(1, silence_duration)],
        dim=1
    )
    
    return speech

# 加载提示音频并完全匹配官方处理
def load_and_process_prompt(file_path):
    """
    严格按照官方demo的提示音频处理流程
    """
    # 使用torchaudio加载保持原始精度
    speech, sr = torchaudio.load(file_path)
    
    # 转为单声道并重采样到模型需要的16kHz (注意! 官方零样本模型需要16kHz提示)
    if sr != 16000:
        speech = torchaudio.transforms.Resample(
            orig_freq=sr, new_freq=16000
        )(speech)
    
    # 应用官方处理流程
    return official_audio_processing(speech, sample_rate=16000)

# 流式合成与播放 (使用官方推荐参数)
def stream_and_play(text, prompt_text, prompt_audio):
    """
    完全匹配官方demo的调用参数
    """
    # 设置随机种子确保结果可复现
    seed = random.randint(1, 100000)
    torch.manual_seed(seed)
    
    print(f"使用随机种子: {seed}")
    print(f"开始合成: {text[:30]}...")
    
    # 使用模型推荐的调用参数
    generator = cosyvoice.inference_zero_shot(
        text,
        prompt_text,
        prompt_audio,
        stream=True,
        speed=1.0 # 保持标准速度
    )
    
    # 在audio_stream_generator中增加双缓冲
    buffer = queue.Queue(maxsize=3)
    def _fill_buffer():
        for chunk in generator:
            audio_tensor = chunk['tts_speech']
            
            # 1. 应用官方后处理流程
            processed_audio = official_audio_processing(audio_tensor, SAMPLE_RATE)
            
            # 2. 转换为float32 numpy数组
            audio_np = processed_audio.squeeze().cpu().numpy().astype(np.float32)
            
            # 3. 实时播放
            # stream.write(audio_np.tobytes())
            buffer.put(audio_np.tobytes())
    
    threading.Thread(target=_fill_buffer, daemon=True).start()
    print("合成完成")
    # 3. 实时播放
    while True:
        stream.write(buffer.get())

# 创建UI类
class TTSApp:
    def __init__(self, root):
        self.root = root
        self.root.title("CosyVoice TTS 合成器")
        self.root.geometry("800x600")
        
        # 设置固定提示文本和音频
        self.prompt_text = "你不要让我做饭啦，我什么都能办到，但是真的不会做饭。"
        self.processed_prompt = load_and_process_prompt('asset/ying.mp3')
        print(f"提示音频已加载: shape={self.processed_prompt.shape}")
        
        # 创建UI元素
        self.create_widgets()
    
    def create_widgets(self):
        # 输入标签
        Label(self.root, text="输入要合成的文本:").pack(pady=5)
        
        # 文本输入框（滚动文本框）
        self.text_input = scrolledtext.ScrolledText(
            self.root, 
            wrap=tk.WORD,
            height=8,
            font=("Arial", 12)
        )
        self.text_input.pack(pady=10, padx=20, fill=tk.BOTH, expand=True)
        self.text_input.insert("1.0", "我真的要谢谢你啊，你真是个天才！")
        
        # 生成按钮
        self.generate_btn = Button(
            self.root,
            text="生成语音",
            command=self.on_generate,
            font=("Arial", 12),
            bg="#4CAF50",
            fg="white",
            height=2,
            width=15
        )
        self.generate_btn.pack(pady=15)
        
        # 状态标签
        self.status_label = Label(
            self.root, 
            text="准备就绪",
            font=("Arial", 10),
            fg="gray"
        )
        self.status_label.pack(pady=5)
        
        # 说明标签
        Label(self.root, text="提示: 生成过程可能需要一些时间，请耐心等待").pack(pady=5)
        Label(self.root, text="提示音频内容: " + self.prompt_text[:50] + "...").pack(pady=5)
    
    def on_generate(self):
        # 获取输入文本
        text = self.text_input.get("1.0", "end-1c").strip()
        
        if not text:
            self.status_label.config(text="错误: 请输入文本", fg="red")
            return
        
        # 更新状态
        self.status_label.config(text="生成中...", fg="blue")
        # self.generate_btn.config(state=tk.DISABLED)
        self.root.update()
        
        # 在新线程中执行生成任务
        threading.Thread(
            target=self.generate_tts, 
            args=(text,),
            daemon=True
        ).start()
    
    def generate_tts(self, text):
        try:
            stream_and_play(text, self.prompt_text, self.processed_prompt)
            self.status_label.config(text="完成! 准备生成下一段文本", fg="green")
        except Exception as e:
            self.status_label.config(text=f"错误: {str(e)}", fg="red")
        finally:
            self.generate_btn.config(state=tk.NORMAL)
            self.root.focus_set()

# 程序主入口
if __name__ == "__main__":
    # 创建GUI
    root = tk.Tk()
    app = TTSApp(root)
    
    # 在程序退出时清理资源
    def on_close():
        stream.stop_stream()
        stream.close()
        p.terminate()
        root.destroy()
    
    root.protocol("WM_DELETE_WINDOW", on_close)
    root.mainloop()